import os
from typing import List
from pathlib import Path
from pydantic import SecretStr
from datetime import datetime
from uuid import uuid4
import logging
from contextlib import asynccontextmanager

# fastapi 相关导入
from fastapi import FastAPI, File, UploadFile, HTTPException, status, Request
from fastapi.responses import JSONResponse
import uvicorn

# RAG的Langchain相关导入
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_milvus import Milvus
from pymilvus import DataType
from langchain_community.document_loaders import (
    PyPDFLoader,
    TextLoader,
    Docx2txtLoader
)

# agent相关导入
from langchain_core.messages import HumanMessage, BaseMessage
from langgraph.prebuilt import create_react_agent
from langchain_core.runnables import RunnableConfig
from langgraph.checkpoint.memory import InMemorySaver
from langchain_deepseek import ChatDeepSeek
from langchain_core.tools import tool

from utils import *

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

# 全局变量存储向量存储实例
vector_store = None

# 初始化RAG组件
embedding_model = HuggingFaceEmbeddings(
    model_name=EMBEDDING_MODEL,
    # model_kwargs={"device": "cpu"}  # 强制使用cpu进行嵌入计算
)
logger.info(f"初始化嵌入模型: {EMBEDDING_MODEL}")

@asynccontextmanager
async def lifespan(app: FastAPI):
    """
    应用启动和关闭时的生命周期管理器
    """
    global vector_store
    # 应用启动时执行的初始化代码
    logger.info("应用正在启动...")
    vector_store = init_milvus()
    os.makedirs(UPLOAD_DIR, exist_ok=True)
    os.makedirs(VECTOR_DB_DIR, exist_ok=True)
    
    yield  # 应用运行期间
    
    # 应用关闭时执行的清理代码
    logger.info("应用正在关闭...")


# 实例化FastAPI
app = FastAPI(lifespan=lifespan)

# 文档加载器映射
LOADERS = {
    ".pdf": PyPDFLoader,
    ".txt": TextLoader,
    ".docx": Docx2txtLoader
}

# 初始化内存存储器
memory = InMemorySaver()

# 添加请求ID中间件以增强日志追踪
@app.middleware("http")
async def add_request_id(request: Request, call_next):
    request_id = str(uuid4())
    request.state.request_id = request_id
    response = await call_next(request)
    return response


# 添加全局异常处理器
@app.exception_handler(Exception)
async def global_exception_handler(request: Request, exc: Exception):
    logger.error(f"未处理的异常: {str(exc)}", exc_info=True)
    return JSONResponse(
        status_code=500,
        content={"detail": "内部服务器错误"}
    )


# 初始化Milvus向量存储
def init_milvus():
    """
    初始化Milvus连接
    """
    try:
        connection_args = {
            "host": MILVUS_HOST,
            "port": MILVUS_PORT,
        }
        
        # 如果提供了用户名和密码，则添加认证信息
        if MILVUS_USERNAME and MILVUS_PASSWORD:
            connection_args["user"] = MILVUS_USERNAME
            connection_args["password"] = MILVUS_PASSWORD
            
        logger.info(f"正在连接到Milvus服务器: {MILVUS_HOST}:{MILVUS_PORT}")
        
        metadata_schema = {
                "filename": {"dtype": DataType.VARCHAR, "max_length": 65535},
                "chunk_index": {"dtype": DataType.INT64},
                "moddate": {"dtype": DataType.VARCHAR, "max_length": 65535},
                "source": {"dtype": DataType.VARCHAR, "max_length": 65535},
                "creationdate": {"dtype": DataType.VARCHAR, "max_length": 65535},
                "page": {"dtype": DataType.INT64},
                "total_pages": {"dtype": DataType.INT64},
            }

        # 创建Milvus实例，参数与现有集合匹配
        vector_store = Milvus(
            embedding_function=embedding_model,
            collection_name=MILVUS_COLLECTION_NAME,
            connection_args=connection_args,
            # 匹配现有集合的关键配置
            enable_dynamic_field=False,  # 现有集合使用静态字段
            auto_id=True,                # 根据主键类型设置
            drop_old=False,               # 是否保留现有数据
            primary_field="pk",           # 主键字段
            text_field="text",            # 文本字段
            vector_field="vector",        # 向量字段
            # 定义元数据字段 schema 以匹配现有集合
            metadata_schema=metadata_schema
        )
        logger.info(f"成功初始化Milvus连接，使用集合: {MILVUS_COLLECTION_NAME}")
        return vector_store
    except Exception as e:
        logger.error(f"初始化Milvus连接失败: {str(e)}")
        return None


# 定义RAG工具用于查询文档
@tool
def query_documents(query: str) -> str:
    """
    查询已上传的文档内容
    
    Args:
        query: 用户的查询问题
        
    Returns:
        str: 检索到的相关文档内容
    """
    global vector_store
    logger.info(f"开始查询文档: {query}")
    
    if vector_store is None:
        logger.warning("向量数据库未初始化")
        return "向量数据库未初始化，请联系管理员。"
    
    try:
        # 检索相关文档
        retriever = vector_store.as_retriever(search_kwargs={"k": 4})
        relevant_docs = retriever.get_relevant_documents(query)
        logger.info(f"检索到 {len(relevant_docs)} 个相关文档")
        
        if not relevant_docs:
            logger.info("未找到与问题相关的内容")
            return "未找到与问题相关的内容。"
        
        # 格式化文档内容
        formatted_docs = "\n\n".join(
            [f"文档来源: {doc.metadata.get('filename', 'unknown')}\n内容: {doc.page_content}" 
             for doc in relevant_docs]
        )
        
        logger.info("成功格式化检索结果")
        return formatted_docs
    except Exception as e:
        logger.error(f"查询文档时出错: {str(e)}")
        return f"查询文档时出错: {str(e)}"

def get_file_loader(filename: str):
    """
    根据文件类型获取相应的加载器
    """
    extension = Path(filename).suffix.lower()
    loader = LOADERS.get(extension)
    if loader:
        logger.debug(f"为文件 {filename} 找到加载器: {loader.__name__}")
    else:
        logger.warning(f"未找到文件 {filename} 的加载器，扩展名: {extension}")
    return loader

@app.get("/")
async def root():
    """
    根端点，提供API信息
    """
    logger.info("访问根端点")
    return JSONResponse(content={
        "message": "文档RAG系统API正在运行", 
        "version": "1.0.0",
        "endpoints": {
            "upload_file": "POST /upload",
            "upload_multiple_files": "POST /upload/multiple",
            "list_files": "GET /files",
            "delete_file": "DELETE /files/{filename}",
            "query": "POST /query"
        }
    })

async def _process_single_file(file: UploadFile):
    """
    处理单个文件上传的内部函数
    """
    logger.info(f"处理文件: {file.filename}")
    # 检查文件是否有名称
    if file.filename is None:
        logger.warning("文件没有文件名")
        return {"filename": "未知", "error": "文件名不能为空"}
    
    # 清理文件名
    secure_name = secure_filename(file.filename)
    if not secure_name:
        logger.warning(f"无效的文件名: {file.filename}")
        return {"filename": file.filename, "error": "无效的文件名"}
    
    # 检查文件扩展名
    if not is_allowed_file(file.filename):
        file_extension = Path(file.filename).suffix.lower()
        logger.warning(f"不允许的文件类型: {file_extension}")
        return {"filename": file.filename, "error": f"文件类型 {file_extension} 不被允许"}
    
    # 保存文件到磁盘
    file_location = os.path.join(UPLOAD_DIR, file.filename)
    logger.info(f"文件将保存到: {file_location}")
    
    # 检查文件是否已存在
    if os.path.exists(file_location):
        logger.warning(f"文件已存在: {file.filename}")
        return {"filename": file.filename, "error": "文件已存在"}
    
    # 保存文件
    try:
        size = 0
        with open(file_location, "wb+") as file_object:
            while True:
                chunk = await file.read(1024*1024)  # 读取1MB块
                if not chunk:
                    break
                size += len(chunk)
                if size > MAX_FILE_SIZE:
                    # 移除部分上传的文件
                    os.remove(file_location)
                    logger.warning(f"文件大小超过限制: {size} > {MAX_FILE_SIZE}")
                    return {"filename": file.filename, "error": f"文件大小超过最大允许大小"}
                file_object.write(chunk)
        
        # 为RAG处理文件
        await process_file_for_rag(file_location, file.filename)
        logger.info(f"文件 {file.filename} 处理完成")
        
        return {
            "filename": file.filename,
            "content_type": file.content_type,
            "size": size
        }
    except Exception as e:
        # 出现任何错误时清理
        if os.path.exists(file_location):
            os.remove(file_location)
            logger.info(f"清理已上传的文件: {file_location}")
        logger.error(f"保存文件失败: {str(e)}")
        return {"filename": file.filename, "error": str(e)}


@app.post("/upload")
async def upload_file(file: UploadFile = File(...)):
    """
    上传单个文件 (PDF, DOCX, 或 TXT) 并为RAG处理
    """
    logger.info(f"开始上传文件: {file.filename}")
    
    # 检查文件是否有名称
    if file.filename is None:
        logger.warning("上传的文件没有文件名")
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail="文件名不能为空"
        )
    
    # 清理文件名
    secure_name = secure_filename(file.filename)
    if not secure_name:
        logger.warning(f"无效的文件名: {file.filename}")
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail="无效的文件名"
        )
    
    # 检查文件扩展名
    if not is_allowed_file(file.filename):
        file_extension = Path(file.filename).suffix.lower()
        logger.warning(f"不允许的文件类型: {file_extension}")
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail=f"文件类型 {file_extension} 不被允许。只允许上传 PDF, DOCX, 和 TXT 文件。"
        )
    
    # 保存文件到磁盘
    file_location = os.path.join(UPLOAD_DIR, file.filename)
    logger.info(f"文件将保存到: {file_location}")
    
    # 检查文件是否已存在
    if os.path.exists(file_location):
        logger.warning(f"文件已存在: {file.filename}")
        raise HTTPException(
            status_code=status.HTTP_409_CONFLICT,
            detail=f"文件 {file.filename} 已存在"
        )
    
    # 保存文件
    try:
        with open(file_location, "wb+") as file_object:
            # 限制文件大小
            size = 0
            while True:
                chunk = await file.read(1024*1024)  # 读取1MB块
                if not chunk:
                    break
                size += len(chunk)
                if size > MAX_FILE_SIZE:
                    # 移除部分上传的文件
                    os.remove(file_location)
                    logger.warning(f"文件大小超过限制: {size} > {MAX_FILE_SIZE}")
                    raise HTTPException(
                        status_code=status.HTTP_413_REQUEST_ENTITY_TOO_LARGE,
                        detail=f"文件大小超过最大允许大小 {MAX_FILE_SIZE} 字节"
                    )
                file_object.write(chunk)
                
        logger.info(f"文件保存成功，大小: {size} 字节")
        # 为RAG处理文件
        await process_file_for_rag(file_location, file.filename)
        logger.info(f"文件 {file.filename} 处理完成")
                
    except Exception as e:
        # 出现任何错误时清理
        if os.path.exists(file_location):
            os.remove(file_location)
            logger.info(f"清理已上传的文件: {file_location}")
        logger.error(f"保存文件失败: {str(e)}")
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail=f"保存文件失败: {str(e)}"
        )
    
    return JSONResponse(content={
        "filename": file.filename,
        "content_type": file.content_type,
        "size": size,
        "message": "文件上传并处理成功"
    })

@app.post("/upload/multiple")
async def upload_multiple_files(files: List[UploadFile] = File(...)):
    """
    上传多个文件 (PDF, DOCX, 或 TXT) 并为RAG处理
    """
    logger.info(f"开始上传 {len(files)} 个文件")
    uploaded_files = []
    errors = []
    
    for file in files:
        try:
            # 创建一个内部函数来处理单个文件上传
            result = await _process_single_file(file)
            if "error" in result:
                errors.append({"filename": result.get("filename", "未知"), "error": result["error"]})
            else:
                uploaded_files.append(result)
        except Exception as e:
            logger.error(f"处理文件时出错: {str(e)}")
            filename = file.filename if file.filename else "未知"
            errors.append({"filename": filename, "error": str(e)})
    
    logger.info(f"批量上传完成: {len(uploaded_files)} 个成功, {len(errors)} 个错误")
    return JSONResponse(content={
        "uploaded_files": uploaded_files,
        "errors": errors,
        "message": f"成功上传并处理 {len(uploaded_files)} 个文件，{len(errors)} 个错误"
    })


async def process_file_for_rag(file_path: str, filename: str):
    """
    为RAG系统处理上传的文件
    """
    logger.info(f"开始处理文件用于RAG: {filename}")
    try:
        # 根据文件类型获取相应的加载器
        loader_class = get_file_loader(filename)
        if not loader_class:
            raise ValueError(f"不支持的文件类型 {filename}")
        
        # 加载文档
        logger.info(f"加载文档: {filename}")
        loader = loader_class(file_path)
        documents = loader.load()
        logger.info(f"成功加载 {len(documents)} 个文档")
        
        # 将文档分割成块
        logger.info(f"分割文档，块大小: {CHUNK_SIZE}, 重叠: {CHUNK_OVERLAP}")
        text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=CHUNK_SIZE,
            chunk_overlap=CHUNK_OVERLAP
        )
        splits = text_splitter.split_documents(documents)
        logger.info(f"文档分割完成，共 {len(splits)} 个块")
        
        # 获取文件最后修改时间
        mod_time = os.path.getmtime(file_path)
        mod_date = datetime.fromtimestamp(mod_time).isoformat()
        creation_time = datetime.fromtimestamp(os.path.getctime(file_path)).isoformat()

        # 添加元数据到每个文档块
        for i, split in enumerate(splits):
            # 确保所有必需的字段都被填充
            split.metadata["filename"] = filename
            split.metadata["chunk_index"] = i
            split.metadata["moddate"] = mod_date
            split.metadata["source"] = file_path
            split.metadata["creationdate"] = creation_time
            split.metadata["page"] = split.metadata.get("page", 0)
            split.metadata["total_pages"] = split.metadata.get("total_pages", 0)
            
            # 确保所有字段都存在，即使值为空
            required_fields = [
                "filename", "chunk_index", "moddate", "source", 
                "creationdate", "page", "total_pages"
            ]
            
            for field in required_fields:
                if field not in split.metadata:
                    split.metadata[field] = ""
        
        # 存储到向量数据库 (Milvus)
        if vector_store is None:
            logger.error("Milvus连接未初始化")
            raise Exception("Milvus连接未初始化")
        
        # 分批添加文档
        batch_size = 5
        for i in range(0, len(splits), batch_size):
            batch = splits[i:i+batch_size]
            vector_store.add_documents(batch)
            del batch
        
        # vector_store.add_documents(splits)
        logger.info(f"成功将文档 {filename} 添加到向量数据库")
        return True
        
    except Exception as e:
        logger.error(f"处理文件用于RAG时出错: {str(e)}")
        raise Exception(f"处理文件用于RAG时出错: {str(e)}")

@app.get("/files")
async def list_files():
    """
    列出所有已上传的文件
    """
    logger.info("列出所有上传的文件")
    files = []
    for filename in os.listdir(UPLOAD_DIR):
        file_path = os.path.join(UPLOAD_DIR, filename)
        if os.path.isfile(file_path):
            file_size = os.path.getsize(file_path)
            files.append({
                "filename": filename,
                "size": file_size,
                "size_formatted": f"{file_size / 1024:.2f} KB" if file_size < 1024*1024 else f"{file_size / (1024*1024):.2f} MB"
            })
    logger.info(f"找到 {len(files)} 个文件")
    return JSONResponse(content={"files": files, "count": len(files)})

@app.delete("/files/{filename}")
async def delete_file(filename: str):
    """
    删除指定文件
    """
    logger.info(f"删除文件: {filename}")
    # 清理文件名以防止目录遍历
    file_path = os.path.join(UPLOAD_DIR, filename)
    
    if not os.path.exists(file_path):
        logger.warning(f"尝试删除不存在的文件: {filename}")
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND,
            detail=f"文件 {filename} 未找到"
        )
    
    try:
        os.remove(file_path)
        logger.info(f"文件删除成功: {filename}")
        return JSONResponse(content={"message": f"文件 {filename} 删除成功"})
    except Exception as e:
        logger.error(f"删除文件失败: {str(e)}")
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail=f"删除文件失败: {str(e)}"
        )

@app.post("/query")
async def query_rag_system(request: Request, query: dict):
    """
    使用问题查询RAG系统
    """
    question = query.get("question")
    logger.info(f"收到查询请求: {question}")
    
    if not question:
        logger.warning("查询问题为空")
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail="问题不能为空"
        )
    
    # 获取用户会话ID，如果不存在则生成新的
    # 可以通过请求头或查询参数传递session_id
    session_id = request.headers.get("X-Session-ID") or query.get("session_id")
    if not session_id:
        # 为新用户生成唯一会话ID
        session_id = str(uuid4())
        logger.info(f"为新会话生成ID: {session_id}")
    else:
        logger.info(f"使用现有会话ID: {session_id}")
    
    try:
        # 检查是否配置了DeepSeek API密钥
        if not DEEPSEEK_API_KEY or DEEPSEEK_API_KEY == "your_api_key_here":
            logger.warning("DeepSeek API密钥未配置")
            return JSONResponse(content={
                "question": question,
                "answer": "DeepSeek API密钥未配置，请在.env文件中设置DEEPSEEK_API_KEY。",
                "session_id": session_id
            })
        
        # 初始化DeepSeek模型
        try:
            logger.info(f"初始化DeepSeek模型: {DEEPSEEK_MODEL}")
            model = ChatDeepSeek(
                model=DEEPSEEK_MODEL,
                api_key=SecretStr(DEEPSEEK_API_KEY),
                api_base=DEEPSEEK_API_BASE
            )
            logger.info("DeepSeek模型初始化成功")
        except Exception as e:
            logger.error(f"初始化DeepSeek模型失败: {str(e)}")
            return JSONResponse(content={
                "question": question,
                "answer": f"初始化DeepSeek模型失败: {str(e)}",
                "session_id": session_id
            })
        
        # 创建ReAct Agent
        try:
            logger.info("创建ReAct Agent")
            agent = create_react_agent(
                model, 
                [query_documents],
                checkpointer=memory
            )
            logger.info("ReAct Agent创建成功")
        except Exception as e:
            logger.error(f"创建ReAct Agent失败: {str(e)}")
            return JSONResponse(content={
                "question": question,
                "answer": f"创建ReAct Agent失败: {str(e)}",
                "session_id": session_id
            })
        
        # 执行查询
        try:
            logger.info("执行查询")
            human_message = HumanMessage(content=question)
            inputs = {"messages": [human_message]}
            
            # 为每个用户设置独立的配置，确保对话隔离
            config = RunnableConfig({"configurable": {"thread_id": session_id}})
            
            # 执行Agent
            result = await agent.ainvoke(inputs, config=config)
            
            # 从结果中提取最后一条消息作为答案
            messages: List[BaseMessage] = result["messages"]
            last_message = messages[-1] if messages else None
            answer = last_message.content if last_message else "未能生成答案。"
            logger.info("查询执行成功")
            
            return JSONResponse(content={
                "question": question,
                "answer": answer,
                "session_id": session_id
            })
        except Exception as e:
            logger.error(f"Agent执行失败: {str(e)}")
            # 如果Agent执行失败，尝试直接使用工具查询
            try:
                logger.info("尝试直接查询文档作为备选方案")
                docs_result = query_documents.invoke(question)
                logger.info("备选方案执行成功")
                return JSONResponse(content={
                    "question": question,
                    "answer": f"基于检索到的文档内容:\n\n{docs_result}\n\n（注意：由于模型配置问题，未使用推理生成答案）",
                    "session_id": session_id
                })
            except Exception as fallback_e:
                logger.error(f"备选方案也失败: {str(fallback_e)}")
                return JSONResponse(content={
                    "question": question,
                    "answer": f"查询过程中出现错误: {str(e)}\n备选方案也失败: {str(fallback_e)}",
                    "session_id": session_id
                })
        
    except Exception as e:
        logger.error(f"查询RAG系统时出错: {str(e)}")
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail=f"查询RAG系统时出错: {str(e)}"
        )

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)